#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>

#define MAX_VIDEO_FRAME_SIZE 400 * 1024
#define RTP_HEADER_SIZE 30
#define RTP_BODY_SIZE 950

int main(int argc, char const *argv[])
{
    int ret;
    AVIOContext *server, *client;
    AVDictionary *in_opts = NULL;

    // ========================= onstruct output rtmp stream start ================================
    AVDictionary *out_opts = NULL;
    av_dict_set(&out_opts, "listen", "2", 0);
    // AVIOContext *avio_out_server_ctx;
    // avio_open2(&avio_out_server_ctx, "rtmp://0.0.0.0:1936/stream/live", AVIO_FLAG_WRITE, NULL, &out_opts);

    AVFormatContext *out_server_ctx = NULL;
    avformat_alloc_output_context2(&out_server_ctx, NULL, "flv", "rtmp://0.0.0.0:1936/stream/live");
    ret = avio_open2(&out_server_ctx->pb, "rtmp://0.0.0.0:1936/stream/live", AVIO_FLAG_WRITE, NULL, &out_opts);
    if (ret < 0)
    {
        av_log(out_server_ctx, AV_LOG_ERROR, "open output stream error");
        return 2;
    }

    // out_server_ctx->pb = avio_out_server_ctx;
    // avformat_find_stream_info(out_server_ctx, NULL);
    av_dump_format(out_server_ctx, 0, NULL, 1);
    // rtmp stream prepare end

    // av stream construct
    AVStream *out_stream;
    out_stream = avformat_new_stream(out_server_ctx, NULL);
    if (!out_stream)
    {
        fprintf(stderr, "Failed allocating output stream\n");
        ret = AVERROR_UNKNOWN;
    }
    // out_stream->time_base.den = 1;
    // out_stream->time_base.num = 25;

    out_stream->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
    out_stream->codecpar->codec_id = AV_CODEC_ID_FLV1;
    out_stream->codecpar->height = 360;
    out_stream->codecpar->width = 360;

    if (av_stream_new_side_data(out_stream, AV_PKT_DATA_NEW_EXTRADATA, 1) < 0)
    {
        av_log(out_server_ctx, AV_LOG_ERROR, "add side data error\n");
        return 4;
    }

    // out_stream->codecpar->codec_tag = 0;
    if (avformat_write_header(out_server_ctx, NULL) < 0)
    {
        av_log(out_server_ctx, AV_LOG_ERROR, "write header error\n");
        return 3;
    }

    // // ========================= onstruct output rtmp stream end ================================

    // ========================== start receive adas device ==========================
    // av_dict_set(&in_opts, "timeout", "2", 0);
    av_dict_set(&in_opts, "listen", "2", 0);

    ret = avio_open2(&server, "tcp://0.0.0.0:9762", AVIO_FLAG_READ, NULL, &in_opts);
    while (1)
    {
        avio_accept(server, &client);

        time_t now;
        time(&now);

        char *file_name;
        // snprintf(file_name, 64, "/home/wzp/h264/video_%lu.h264", now);
        file_name = "/home/wzp/h264/video_2.h264";

        FILE *fp;
        fp = fopen(file_name, "wb+");

        uint8_t *header_buf = av_mallocz(RTP_HEADER_SIZE);
        uint8_t *body_buf = av_mallocz(RTP_BODY_SIZE);
        uint8_t *video_buf = av_mallocz(MAX_VIDEO_FRAME_SIZE);

        uint32_t video_buf_write_idx = 0;

        int pts = 0;
        int dts = 0;
        while (1)
        {
            ret = avio_read(client, header_buf, RTP_HEADER_SIZE);
            if (ret < 0)
            {
                av_free(header_buf);
                av_free(body_buf);
                fclose(fp);
                break;
            }

            if (header_buf[0] == 0x30 && header_buf[1] == 0x31 && header_buf[2] == 0x63 && header_buf[3] == 0x64)
            {
                if (!video_buf)
                {
                    video_buf = av_mallocz(MAX_VIDEO_FRAME_SIZE);
                }
                uint8_t is_complete_and_mtype = header_buf[5];
                uint16_t tcp_body_size = (*(header_buf + 28) << 8) | *(header_buf + 29);

                uint16_t last_frame_time_gap = (*(header_buf + 26) << 8) | *(header_buf + 27);

                ret = avio_read(client, body_buf, tcp_body_size);
                if (ret < 0)
                {
                    av_free(header_buf);
                    av_free(body_buf);
                    fclose(fp);
                    break;
                }

                // if it is h264 data, copy and process
                if ((is_complete_and_mtype & 0x7f) == 0x62)
                {
                    memcpy(video_buf + video_buf_write_idx, body_buf, tcp_body_size);
                    video_buf_write_idx += tcp_body_size;

                    // is last packet
                    if ((is_complete_and_mtype >> 7) == 1)
                    {
                        ret = fwrite(video_buf, 1, video_buf_write_idx, fp);
                        if (ret < 0)
                        {
                            printf("file write error\n");
                            break;
                        }

                        AVPacket *pkt = av_packet_alloc();
                        // send packet
                        ret = av_packet_from_data(pkt, video_buf, video_buf_write_idx);

                        if (ret < 0)
                        {
                            break;
                        }
                        // pts, dts, stream_idx, duration, time_base
                        pkt->stream_index = 0;
                        pkt->duration = last_frame_time_gap;
                        pkt->pts = ++pts;
                        pkt->dts = ++dts;
                        pkt->pos += video_buf_write_idx;
                        pkt->flags = 1;

                        AVPacket *send_out_packet = av_packet_clone(pkt);

                        ret = av_interleaved_write_frame(out_server_ctx, send_out_packet);
                        if (ret < 0)
                        {
                            av_log(out_server_ctx, AV_LOG_ERROR, "interleaved write error%d\n", ret);
                            break;
                        }

                        printf("whole video packet size:%d\n", video_buf_write_idx);
                        // memset(video_buf, 0, video_buf_write_idx);
                        video_buf_write_idx = 0;
                    }
                }
                else
                {
                    printf("audio data\n");
                }
            }
            else
            {
                printf("data error, exit tcp channel\n");
                break;
            }

            // clear buf
            memset(header_buf, 0, RTP_HEADER_SIZE);
            memset(body_buf, 0, RTP_BODY_SIZE);
        }

        av_free(header_buf);
        av_free(body_buf);
        fclose(fp);
    }

    return 0;
}
